package com.mahiru.mahiruaiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.mahiru.mahiruaiagent.agent.model.AgentState;
import com.mahiru.mahiruaiagent.common.ErrorCode;
import com.mahiru.mahiruaiagent.exception.BusinessException;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * 抽象基础代理类 管理代理状态和执行流程
 *
 * @Author Enkidu
 * @Date 2025/7/9 16:33
 */
@Data
@Slf4j
public abstract class BaseAgent {

    // 核心属性
    private String name;

    // 系统提示词
    private String systemPrompt;
    // 中间提示词
    private String nextStepPrompt;

    // 代理状态
    private AgentState state = AgentState.IDLE;

    // 执行步骤控制
    private int currentStep = 0;
    private int maxSteps = 10;

    // 大模型
    private ChatClient chatClient;

    // Memory 记忆 （自主维护会话上下文记忆）
    private List<Message> messageList = new ArrayList<>();

    /**
     * 执行代理的主流程
     *
     * @param userPrompt 用户提示词
     * @return 运行结果
     */

    public String run(String userPrompt) {
        // ThrowUtils.throwIf(StrUtil.isBlank(userPrompt), ErrorCode.PARAMS_ERROR);
        if (StrUtil.isBlank(userPrompt)) {
            log.warn("User prompt is blank");
            throw new BusinessException(ErrorCode.PARAMS_ERROR, "User prompt cannot be blank");
        }
        if (state != AgentState.IDLE) {
            log.warn("Agent {} is not idle, current state: {}", name, state);
            throw new BusinessException(ErrorCode.OPERATION_ERROR, "Cannot run agent, current state: " + state);
        }
        // 执行
        state = AgentState.RUNNING;
        // 记录消息上下文
        messageList.add(new UserMessage(userPrompt));
        // 保存结果列表
        List<String> results = new ArrayList<>();
        try {
            // 执行循环
            for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Executing step {}/{} for agent: {}", stepNumber, maxSteps, name);
                // 单步执行
                String stepResult = step();
                String result = "Step" + stepNumber + ": " + stepResult;
                results.add(result);
            }
            // 检查是否超出步骤
            if (currentStep >= maxSteps) {
                state = AgentState.FINISHED;
                results.add("Reached maximum steps: " + maxSteps);
                log.warn("Reached maximum steps for agent: {}, current step: {}", name, currentStep);
            }
            return String.join("\n", results);
        } catch (Exception e) {
            state = AgentState.ERROR;
            log.error("Error occurred during agent execution: {}", e.getMessage(), e);
            return "执行错误: " + e.getMessage();
        } finally {
            // 清理资源
            cleanup();
        }
    }

    /**
     * 执行代理 (流式执行)
     *
     * @param userPrompt 用户提示词
     * @return 运行结果
     */

    public SseEmitter runStream(String userPrompt) {
        SseEmitter sseEmitter = new SseEmitter(300000L);
        // 异步执行
        CompletableFuture.runAsync(() -> {
// ThrowUtils.throwIf(StrUtil.isBlank(userPrompt), ErrorCode.PARAMS_ERROR);
            try {
                if (StrUtil.isBlank(userPrompt)) {
                    log.warn("User prompt is blank");
                    sseEmitter.send("不能使用空提示词！");
                    sseEmitter.complete();
                    // throw new BusinessException(ErrorCode.PARAMS_ERROR, "User prompt cannot be blank");
                    return;
                }
                if (state != AgentState.IDLE) {
                    log.warn("Agent {} is not idle, current state: {}", name, state);
                    sseEmitter.send("代理当前状态不允许执行: " + state);
                    sseEmitter.complete();
                    // throw new BusinessException(ErrorCode.OPERATION_ERROR, "Cannot run agent, current state: " + state);
                    return;
                }
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }

            // 执行
            state = AgentState.RUNNING;
            // 记录消息上下文
            messageList.add(new UserMessage(userPrompt));
            // 保存结果列表
            List<String> results = new ArrayList<>();
            try {
                // 执行循环
                for (int i = 0; i < maxSteps && state != AgentState.FINISHED; i++) {
                    int stepNumber = i + 1;
                    currentStep = stepNumber;
                    log.info("Executing step {}/{} for agent: {}", stepNumber, maxSteps, name);
                    // 单步执行
                    String stepResult = step();
                    String result = "Step" + stepNumber + ": " + stepResult;
                    results.add(result);
                    // 发送当前步骤结果
                    sseEmitter.send(result);
                }
                // 检查是否超出步骤
                if (currentStep >= maxSteps) {
                    state = AgentState.FINISHED;
                    results.add("Reached maximum steps: " + maxSteps);
                    log.warn("Reached maximum steps for agent: {}, current step: {}", name, currentStep);
                    sseEmitter.send("已达到最大步骤: " + maxSteps + "，执行结束！");
                }
                // 正常完成
                sseEmitter.complete();
            } catch (Exception e) {
                state = AgentState.ERROR;
                log.error("Error occurred during agent execution: {}", e.getMessage(), e);
                try {
                    sseEmitter.send("执行错误: " + e.getMessage());
                    sseEmitter.complete();
                } catch (IOException ex) {
                    sseEmitter.completeWithError(ex);
                }
            } finally {
                // 清理资源
                cleanup();
            }
        });

        // 设置超时回调
        sseEmitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanup();
            log.warn("SSE 连接超时");
        });

        // 正常完成
        sseEmitter.onCompletion(() -> {
            if (this.state == AgentState.RUNNING) {
                this.state = AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE 连接正常完成，代理状态: {}", this.state);
        });
        return sseEmitter;
    }

    /**
     * 执行代理的单步操作
     *
     * @return
     */
    public abstract String step();

    /**
     * 清理资源
     */
    protected void cleanup() {
    }
}
